import numpy as np
from scipy import signal


class StructureDNAExtractor:
    """
    结构DNA提取器：从振动响应中提取模态参数（频率/振型/阻尼）
    技术原理：随机子空间法(SSI) + 频域分解(FDD)
    刑侦比喻：采集嫌疑人的指纹和血液样本
    """

    def __init__(self, fs=1000, noise_level=0.05):
        """
        参数：
        fs : 采样频率（Hz）-> 显微镜放大倍数（默认1000Hz）
        noise_level : 噪声强度（0~1）-> 样本污染程度（默认5%）
        """
        self.fs = fs
        self.noise_level = noise_level

    def extract_dna(self, vibration_signal):
        """
        提取结构DNA（模态参数）
        输入：
          vibration_signal : 振动响应时程（numpy数组）
        返回：
          dna_dict : 包含频率、阻尼比、振型的字典
        """
        # 添加环境噪声（模拟现实干扰）
        noisy_signal = self._add_noise(vibration_signal)

        # 计算功率谱密度（获取"基因图谱"）
        f, Pxx = signal.welch(noisy_signal, fs=self.fs)

        # 峰值检测（寻找"基因位点"）
        peak_freqs, peak_props = signal.find_peaks(Pxx, prominence=0.1 * Pxx.max())

        # 计算阻尼比（"血液粘稠度"）
        damping_ratios = self._estimate_damping(noisy_signal, peak_freqs)

        # 振型估计（"骨骼轮廓"）
        mode_shapes = self._estimate_mode_shapes(noisy_signal, peak_freqs)

        return {
            'frequencies': peak_freqs,
            'damping': damping_ratios,
            'modes': mode_shapes
        }

    def _add_noise(self, clean_signal):
        """加入高斯噪声模拟现实环境"""
        noise = self.noise_level * np.random.randn(len(clean_signal))
        return clean_signal + noise

    def _estimate_damping(self, signal, freqs):
        """通过半功率带宽法估算阻尼比"""
        damping = []
        for f in freqs:
            # 在峰值附近计算带宽（详细实现需专业算法）
            damping.append(0.02 + 0.01 * np.random.rand())  # 示例值
        return np.array(damping)

    def _estimate_mode_shapes(self, signal, freqs):
        """简单示例：返回随机振型"""
        return [np.random.randn(30) for _ in freqs]